import plugin from '../../lib/plugins/plugin.js'
import common from "../../lib/common/common.js";
import fetch from 'node-fetch'
import axios from 'axios'

const API = 'https://ayaka-daisuki-real-cugan.hf.space/api/predict'

let zt = 0
let NumReg = '[一壹二两三四五六七八九十百千万亿\\d]+'
export class example extends plugin {
  constructor() {
    super({
            /** 功能名称 */
            name: '图片清晰术',
            /** 功能描述 */
            dsc: '图片清晰术',
            event: 'message',
            /** 优先级，数字越小等级越高 */
            priority: 1009,
            rule: [
                {
                    /** 命令正则匹配 */
                    reg: '^#?(超分|清晰术)(.*)(降噪)?$',
                    /** 执行方法 */
                    fnc: 'appreciate'
                }
            ]
        })
    }

  async appreciate(e) {
    if (API == '') return 
	if (zt == 1) {
		e.reply('当前有任务正在执行..', true, { recallMsg: 10 })
		return
	}
	let image = await this.getImage(e)
	if (!image) return
	
	let tag = e.msg.replace(/#|(超分|清晰术)|降噪/g, '').trim()
    let num = 2
        let Match = e.msg.match(new RegExp(`(${NumReg})(倍)`));
        if (Match) {
            num = Number(Match[1]);
            if (isNaN(num)) {
                return;
            }
			if (num < 2 || num > 4) {
				e.reply(`参数错误（2倍，3倍，4倍）`, true)
				return
			}
			tag = tag.replace(Match[0], "").trim();  // 去掉匹配的页码部分
        }
	if (!tag) {
		tag = '强力'
	}
	
	    let modelmap = {
		    "conservative": ["保守"],
            "denoise3x": ["强力", "三级", "3级"],
            "no-denoise": ["无", "不"]
		}
        let Model = `up${num}x-latest-denoise3x.pth`
        for (const [model, keywords] of Object.entries(modelmap)) {
            if (keywords.some(keyword => tag.includes(keyword))) {
                Model = `up${num}x-latest-${model}.pth`;
                break;
            }
        }

	    let base64 = await url_to_base64(image)
	    e.reply(`图片超分中，${num}倍 ${tag}降噪...`, true, { recallMsg: 10 })
		
		zt = 1
	    let response = ""
	    let res = ""

	try {
        response = await fetch(API, {
            method: "POST",
            headers: {
                "Content-Type": "application/json"
            },
            body: JSON.stringify({
                data: [
			        "data:image/png;base64," + base64,
			        Model,
			        0,
		        ]
            })
        })
        res = await response.json()
        let msgs = []
        let imagedata = res.data[0].replace(/^data:image\/png;base64,/, "");
        let buffer = Buffer.from(imagedata, 'base64');
		
		msgs.push(`====图片清晰术处理完成====`)
		msgs.push(segment.image(image))
		msgs.push(`========处理结果========`)
		msgs.push(segment.image(buffer))
        //await e.reply(segment.image(buffer), true);
		await e.reply(await common.makeForwardMsg(e, msgs))
		zt = 0
    } catch (error) {
        e.reply('超分失败', true)
		zt = 0
        return
    }

  }
		
  async getImage(e) {
	let imgUrls = []
	  if (e.source || e.reply_id ) {
        // 优先从回复找图
        let reply
        if (this.e.getReply) {
          reply = await this.e.getReply()
        } else if (this.e.source) {
          if (this.e.group?.getChatHistory)
            reply = (await this.e.group.getChatHistory(this.e.source.seq, 1)).pop()
          else if (this.e.friend?.getChatHistory)
            reply = (await this.e.friend.getChatHistory(this.e.source.time, 1)).pop()
        }
        if (reply?.message) {
          for (let val of reply.message) {
            if (val.type === 'image') {
                imgUrls.push(val.url)
            }
          }
        }
      }
	  else if (e.img) {
        // 一起发的图
        imgUrls.push(...e.img)
      }   
	if (imgUrls.length) {
        return imgUrls[0]
	} else {
		return ''
	}
  }
}

async function url_to_base64(url) {
        let img = await axios.get(url, {
            responseType: 'arraybuffer'
        });
        let base64 = Buffer.from(img.data, 'binary')
            .toString('base64');
        return base64
    }